import torch
import clip
from PIL import Image
import numpy as np
import cv2
import matplotlib.pyplot as plt
from io import BytesIO
import requests


def interpret(image, text, model, device, index=None):
    logits_per_image, logits_per_text = model(image, text)
    probs = logits_per_image.softmax(dim=-1).detach().cpu().numpy()
    # print(logits_per_text.shape)
    # print(logits_per_image.shape)
    # probs_text = logits_per_text.softmax(dim=0).detach().cpu().numpy()
    # print(probs_text)
    return probs
    # print(logits_per_text)
    # print(logits_per_image)
    # if index is None:
    #     index = np.argmax(logits_per_image.cpu().data.numpy(), axis=-1)
    # one_hot = np.zeros((1, logits_per_image.size()[-1]), dtype=np.float32)
    # one_hot[0, index] = 1
    # one_hot = torch.from_numpy(one_hot).requires_grad_(True)
    # one_hot = torch.sum(one_hot.cuda() * logits_per_image)
    # model.zero_grad()
    # one_hot.backward(retain_graph=True)
    #
    # image_attn_blocks = list(dict(model.visual.transformer.resblocks.named_children()).values())
    # num_tokens = image_attn_blocks[0].attn_probs.shape[-1]
    # R = torch.eye(num_tokens, num_tokens, dtype=image_attn_blocks[0].attn_probs.dtype).to(device)
    # for blk in image_attn_blocks:
    #     grad = blk.attn_grad
    #     cam = blk.attn_probs
    #     cam = cam.reshape(-1, cam.shape[-1], cam.shape[-1])
    #     grad = grad.reshape(-1, grad.shape[-1], grad.shape[-1])
    #     cam = grad * cam
    #     cam = cam.clamp(min=0).mean(dim=0)
    #     R += torch.matmul(cam, R)
    # R[0, 0] = 0
    # image_relevance = R[0, 1:]
    # print(image_relevance)
    #
    # # create heatmap from mask on image
    # def show_cam_on_image(img, mask):
    #     heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)
    #     heatmap = np.float32(heatmap) / 255
    #     cam = heatmap + np.float32(img)
    #     cam = cam / np.max(cam)
    #     return cam
    #
    # image_relevance = image_relevance.reshape(1, 1, 7, 7)
    # image_relevance = torch.nn.functional.interpolate(image_relevance, size=224, mode='bilinear')
    # image_relevance = image_relevance.reshape(224, 224).cuda().data.cpu().numpy()
    # image_relevance = (image_relevance - image_relevance.min()) / (image_relevance.max() - image_relevance.min())
    # image = image[0].permute(1, 2, 0).data.cpu().numpy()
    # image = (image - image.min()) / (image.max() - image.min())
    # vis = show_cam_on_image(image, image_relevance)
    # vis = np.uint8(255 * vis)
    # vis = cv2.cvtColor(np.array(vis), cv2.COLOR_RGB2BGR)
    #
    # plt.imshow(vis)
    # plt.show()
    #
    # print("Label probs:", probs)

device = "cuda" if torch.cuda.is_available() else "cpu"
clip.available_models()
model, preprocess = clip.load("ViT-B/32", device=device, jit=False)
model.eval()

def calculate_cosine(token, url):
    # 使用 requests 获取图片数据
    response = requests.get(url)
    # 使用 BytesIO 创建一个文件类对象，然后用 Image.open 打开它
    url_image = Image.open(BytesIO(response.content))
    image = preprocess(url_image).unsqueeze(0).to(device)
    text = clip.tokenize(token).to(device)
    return interpret(model=model, image=image, text=text, device=device, index=0)
    # with torch.no_grad():
    #     image_features = model.encode_image(image).float()
    #     text_features = model.encode_text(text).float()
    #
    # text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)
    # print(text_probs)

    # top_probs, top_labels = text_probs.cpu().topk(5, dim=-1)
    #
    # print(top_probs)
    # print(top_labels)
    # interpret(model=model, image=image, text=text, device=device, index=0)
    # interpret(model=model, image=image, text=text, device=device, index=1)


if __name__ == "__main__":
    # 一条狗的图片 https://fakerlove.oss-cn-beijing.aliyuncs.com/kaoyan/2024/03/17/2c7d661889be4fa792ffe6bdd947f1c110.jpg
    #
    # 红色新征程 https://fakerlove.oss-cn-beijing.aliyuncs.com/kaoyan/2024/03/17/b519049322dd401399cf1c7f26c111a08.jpg
    result = calculate_cosine(["a dog", "car"],
                              "https://fakerlove.oss-cn-beijing.aliyuncs.com/kaoyan/2024/03/17/2c7d661889be4fa792ffe6bdd947f1c110.jpg")
    print(result[0][0])
    # print(result[0][1])
    # print(result[0][2])
